#include "SemanticAnalyzer.h"
#include <iostream>
#include <algorithm>

namespace MiniDB {

SemanticAnalyzer::SemanticAnalyzer(SymbolTable& symbol_table)
    : symbol_table_(symbol_table) {}

bool SemanticAnalyzer::analyze(const std::vector<std::shared_ptr<Statement>>& statements) {
    clearErrors();
    bool success = true;

    for (const auto& stmt : statements) {
        if (!analyzeStatement(stmt)) {
            success = false;
        }
    }

    return success;
}

bool SemanticAnalyzer::analyzeStatement(std::shared_ptr<Statement> stmt) {
    if (auto create_stmt = std::dynamic_pointer_cast<CreateTableStatement>(stmt)) {
        return analyzeCreateTable(create_stmt);
    } else if (auto insert_stmt = std::dynamic_pointer_cast<InsertStatement>(stmt)) {
        return analyzeInsert(insert_stmt);
    } else if (auto update_stmt = std::dynamic_pointer_cast<UpdateStatement>(stmt)) {
        return analyzeUpdate(update_stmt);
    } else if (auto delete_stmt = std::dynamic_pointer_cast<DeleteStatement>(stmt)) {
        return analyzeDelete(delete_stmt);
    } else if (auto select_stmt = std::dynamic_pointer_cast<SelectStatement>(stmt)) {
        return analyzeSelect(select_stmt);
    }

    return true; // Unknown statement type, skip
}

bool SemanticAnalyzer::analyzeCreateTable(std::shared_ptr<CreateTableStatement> stmt) {
    // Check if table already exists
    if (symbol_table_.hasTable(stmt->table_name)) {
        addError("TABLE_EXISTS", 0, 0, "Table '" + stmt->table_name + "' already exists");
        return false;
    }

    // Create table info
    TableInfo table_info(stmt->table_name);
    table_info.name = stmt->table_name;

    // Analyze columns
    for (const auto& col_node : stmt->columns) {
        if (auto col_def = std::dynamic_pointer_cast<ColumnDefinition>(col_node)) {
            DataType data_type = stringToDataType(col_def->data_type);
            bool is_pk = std::find(col_def->constraints.begin(), col_def->constraints.end(),
                                 "PRIMARY KEY") != col_def->constraints.end();

            ColumnInfo col_info(col_def->column_name, data_type, 0, true, is_pk);
            table_info.addColumn(col_info);
        }
    }

    // Add table to symbol table
    if (!symbol_table_.addTable(table_info)) {
        addError("TABLE_CREATION_FAILED", 0, 0, "Failed to create table '" + stmt->table_name + "'");
        return false;
    }

    std::cout << "[SEMANTIC] Table '" << stmt->table_name << "' semantic check passed" << std::endl;
    return true;
}

bool SemanticAnalyzer::analyzeInsert(std::shared_ptr<InsertStatement> stmt) {
    // Check if table exists
    if (!checkTableExists(stmt->table_name, 0, 0)) {
        return false;
    }

    const TableInfo* table_info = symbol_table_.getTable(stmt->table_name);
    if (!table_info) return false;

    // Check column count if columns are specified
    if (!stmt->columns.empty() && stmt->values.size() > 0) {
        if (stmt->columns.size() != stmt->values[0].size()) {
            addError("COLUMN_COUNT_MISMATCH", 0, 0,
                    "Column count mismatch: " + std::to_string(stmt->columns.size()) +
                    " columns specified, " + std::to_string(stmt->values[0].size()) + " values provided");
            return false;
        }

        // Check if specified columns exist
        for (const std::string& col_name : stmt->columns) {
            if (!table_info->hasColumn(col_name)) {
                addError("COLUMN_NOT_FOUND", 0, 0,
                        "Column '" + col_name + "' does not exist in table '" + stmt->table_name + "'");
                return false;
            }
        }
    } else if (stmt->columns.empty()) {
        // No columns specified, check value count matches table columns
        if (stmt->values.size() > 0 && stmt->values[0].size() != table_info->columns.size()) {
            addError("COLUMN_COUNT_MISMATCH", 0, 0,
                    "Value count mismatch: expected " + std::to_string(table_info->columns.size()) +
                    " values, got " + std::to_string(stmt->values[0].size()));
            return false;
        }
    }

        // Check data types for each value
    if (stmt->values.size() > 0) {
        const auto& row = stmt->values[0];
        for (size_t i = 0; i < row.size(); ++i) {
            DataType expected_type;
            std::string col_name;

            if (!stmt->columns.empty()) {
                col_name = stmt->columns[i];
            } else {
                col_name = table_info->columns[i].name;
            }

            const ColumnInfo* col_info = table_info->getColumn(col_name);
            if (!col_info) continue;

            expected_type = col_info->type;

            DataType actual_type;
            if (!analyzeExpression(row[i], actual_type)) {
                return false;
            }

            if (!symbol_table_.areTypesCompatible(expected_type, actual_type)) {
                addError("TYPE_MISMATCH", 0, 0,
                        "Type mismatch for column '" + col_name + "': expected " +
                        dataTypeToString(expected_type) + ", got " + dataTypeToString(actual_type));
                return false;
            }
        }
    }

    std::cout << "[SEMANTIC] INSERT into '" << stmt->table_name << "' semantic check passed" << std::endl;
    return true;
}

bool SemanticAnalyzer::analyzeSelect(std::shared_ptr<SelectStatement> stmt) {
    // Check FROM clause
    if (auto from_clause = std::dynamic_pointer_cast<FromClause>(stmt->from_clause)) {
        if (!checkTableExists(from_clause->table_name, 0, 0)) {
            return false;
        }
    }

    // Check WHERE clause
    if (stmt->where_clause) {
        DataType result_type;
        if (!analyzeExpression(stmt->where_clause, result_type)) {
            return false;
        }

        // WHERE clause should evaluate to boolean
        if (result_type != DataType::BOOLEAN) {
            addError("INVALID_WHERE_TYPE", 0, 0,
                    "WHERE clause must evaluate to boolean type, got " + dataTypeToString(result_type));
            return false;
        }
    }

    std::cout << "[SEMANTIC] SELECT statement semantic check passed" << std::endl;
    return true;
}

bool SemanticAnalyzer::analyzeUpdate(std::shared_ptr<UpdateStatement> stmt) {
    if (!checkTableExists(stmt->table_name, 0, 0)) {
        return false;
    }

    const TableInfo* table_info = symbol_table_.getTable(stmt->table_name);
    if (!table_info) return false;

    // Check assignments
    for (const auto& assignment : stmt->assignments) {
        if (!table_info->hasColumn(assignment.first)) {
            addError("COLUMN_NOT_FOUND", 0, 0,
                    "Column '" + assignment.first + "' does not exist in table '" + stmt->table_name + "'");
            return false;
        }

        const ColumnInfo* col_info = table_info->getColumn(assignment.first);
        DataType expr_type;
        if (!analyzeExpression(assignment.second, expr_type)) {
            return false;
        }

        if (!symbol_table_.areTypesCompatible(col_info->type, expr_type)) {
            addError("TYPE_MISMATCH", 0, 0,
                    "Type mismatch for column '" + assignment.first + "': expected " +
                    dataTypeToString(col_info->type) + ", got " + dataTypeToString(expr_type));
            return false;
        }
    }

    // Check WHERE clause
    if (stmt->where_clause) {
        DataType result_type;
        if (!analyzeExpression(stmt->where_clause, result_type)) {
            return false;
        }

        if (result_type != DataType::BOOLEAN) {
            addError("INVALID_WHERE_TYPE", 0, 0,
                    "WHERE clause must evaluate to boolean type");
            return false;
        }
    }

    std::cout << "[SEMANTIC] UPDATE statement semantic check passed" << std::endl;
    return true;
}

bool SemanticAnalyzer::analyzeDelete(std::shared_ptr<DeleteStatement> stmt) {
    if (!checkTableExists(stmt->table_name, 0, 0)) {
        return false;
    }

    // Check WHERE clause
    if (stmt->where_clause) {
        DataType result_type;
        if (!analyzeExpression(stmt->where_clause, result_type)) {
            return false;
        }

        if (result_type != DataType::BOOLEAN) {
            addError("INVALID_WHERE_TYPE", 0, 0,
                    "WHERE clause must evaluate to boolean type");
            return false;
        }
    }

    std::cout << "[SEMANTIC] DELETE statement semantic check passed" << std::endl;
    return true;
}

bool SemanticAnalyzer::analyzeExpression(std::shared_ptr<Expression> expr, DataType& result_type) {
    if (auto binary_expr = std::dynamic_pointer_cast<BinaryExpression>(expr)) {
        return analyzeBinaryExpression(binary_expr, result_type);
    } else if (auto column_ref = std::dynamic_pointer_cast<ColumnReference>(expr)) {
        return analyzeColumnReference(column_ref, result_type);
    } else if (auto literal = std::dynamic_pointer_cast<Literal>(expr)) {
        return analyzeLiteral(literal, result_type);
    }

    result_type = DataType::INT; // Default
    return true;
}

bool SemanticAnalyzer::analyzeBinaryExpression(std::shared_ptr<BinaryExpression> expr, DataType& result_type) {
    DataType left_type, right_type;

    if (!analyzeExpression(expr->left, left_type) ||
        !analyzeExpression(expr->right, right_type)) {
        return false;
    }

    // Comparison operators return boolean
    if (expr->operator_ == "=" || expr->operator_ == "!=" ||
        expr->operator_ == ">" || expr->operator_ == "<" ||
        expr->operator_ == ">=" || expr->operator_ == "<=") {
        result_type = DataType::BOOLEAN;
    } else {
        // Arithmetic operators
        result_type = symbol_table_.getResultType(left_type, right_type);
    }

    return true;
}

bool SemanticAnalyzer::analyzeColumnReference(std::shared_ptr<ColumnReference> expr, DataType& result_type) {
    // For simplicity, assume single table queries
    // In a real system, we'd need to handle table aliases and multi-table queries

    // Try to find column in any table (simplified)
    auto table_names = symbol_table_.getTableNames();
    for (const std::string& table_name : table_names) {
        DataType col_type;
        if (checkColumnExists(table_name, expr->column_name, 0, 0, col_type)) {
            result_type = col_type;
            return true;
        }
    }

    addError("COLUMN_NOT_FOUND", 0, 0, "Column '" + expr->column_name + "' not found in any table");
    return false;
}

bool SemanticAnalyzer::analyzeLiteral(std::shared_ptr<Literal> expr, DataType& result_type) {
    // Determine type based on variant
    if (std::holds_alternative<int>(expr->value)) {
        result_type = DataType::INT;
    } else if (std::holds_alternative<double>(expr->value)) {
        result_type = DataType::FLOAT;
    } else if (std::holds_alternative<std::string>(expr->value)) {
        result_type = DataType::VARCHAR;
    }

    return true;
}

DataType SemanticAnalyzer::stringToDataType(const std::string& type_str) const {
    if (type_str.find("INT") == 0) return DataType::INT;
    if (type_str.find("VARCHAR") == 0) return DataType::VARCHAR;
    if (type_str.find("FLOAT") == 0) return DataType::FLOAT;
    if (type_str.find("BOOLEAN") == 0) return DataType::BOOLEAN;
    if (type_str.find("DATE") == 0) return DataType::DATE;
    if (type_str.find("TEXT") == 0) return DataType::TEXT;
    return DataType::VARCHAR; // Default
}

std::string SemanticAnalyzer::dataTypeToString(DataType type) const {
    switch (type) {
        case DataType::INT: return "INT";
        case DataType::VARCHAR: return "VARCHAR";
        case DataType::FLOAT: return "FLOAT";
        case DataType::BOOLEAN: return "BOOLEAN";
        case DataType::DATE: return "DATE";
        case DataType::TEXT: return "TEXT";
        default: return "UNKNOWN";
    }
}

void SemanticAnalyzer::addError(const std::string& error_type, int line, int column, const std::string& message) {
    errors_.emplace_back(error_type, line, column, message);
    std::cerr << "[SEMANTIC ERROR] " << error_type << " at line " << line
              << ", column " << column << ": " << message << std::endl;
}

bool SemanticAnalyzer::checkTableExists(const std::string& table_name, int line, int column) {
    if (!symbol_table_.hasTable(table_name)) {
        addError("TABLE_NOT_FOUND", line, column, "Table '" + table_name + "' does not exist");
        return false;
    }
    return true;
}

bool SemanticAnalyzer::checkColumnExists(const std::string& table_name, const std::string& column_name,
                                       int line, int column, DataType& column_type) {
    const ColumnInfo* col_info = symbol_table_.getColumn(table_name, column_name);
    if (!col_info) {
        addError("COLUMN_NOT_FOUND", line, column,
                "Column '" + column_name + "' does not exist in table '" + table_name + "'");
        return false;
    }

    column_type = col_info->type;
    return true;
}

} // namespace MiniDB
